// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <fmt/format.h>
#include <chrono>

#include <pollux/common/base/succinct_printer.h>
#include <pollux/common/process/process_base.h>

namespace kumo::pollux {
    // Tracks call count and elapsed CPU and wall time for a repeating operation.
    struct CpuWallTiming {
        uint64_t count = 0;
        uint64_t wallNanos = 0;
        uint64_t cpuNanos = 0;

        void add(const CpuWallTiming &other) {
            count += other.count;
            cpuNanos += other.cpuNanos;
            wallNanos += other.wallNanos;
        }

        void clear() {
            count = 0;
            wallNanos = 0;
            cpuNanos = 0;
        }

        std::string toString() const {
            return fmt::format(
                "count: {}, wallTime: {}, cpuTime: {}",
                count,
                succinctNanos(wallNanos),
                succinctNanos(cpuNanos));
        }
    };

    // Adds elapsed CPU and wall time to a CpuWallTiming.
    class CpuWallTimer {
    public:
        explicit CpuWallTimer(CpuWallTiming &timing);

        ~CpuWallTimer();

    private:
        uint64_t cpuTimeStart_;
        std::chrono::steady_clock::time_point wallTimeStart_;
        CpuWallTiming &timing_;
    };

    /// Keeps track of elapsed CPU and wall time from construction time.
    class DeltaCpuWallTimeStopWatch {
    public:
        explicit DeltaCpuWallTimeStopWatch()
            : wallTimeStart_(std::chrono::steady_clock::now()),
              cpuTimeStart_(process::threadCpuNanos()) {
        }

        CpuWallTiming elapsed() const {
            // NOTE: End the cpu-time timing first, and then end the wall-time timing,
            // so as to avoid the counter-intuitive phenomenon that the final calculated
            // cpu-time is slightly larger than the wall-time.
            uint64_t cpuTimeDuration = process::threadCpuNanos() - cpuTimeStart_;
            uint64_t wallTimeDuration =
                    std::chrono::duration_cast<std::chrono::nanoseconds>(
                        std::chrono::steady_clock::now() - wallTimeStart_)
                    .count();
            return CpuWallTiming{1, wallTimeDuration, cpuTimeDuration};
        }

    private:
        // NOTE: Put `wallTimeStart_` before `cpuTimeStart_`, so that wall-time starts
        // counting earlier than cpu-time.
        const std::chrono::steady_clock::time_point wallTimeStart_;
        const uint64_t cpuTimeStart_;
    };

    /// Composes delta CpuWallTiming upon destruction and passes it to the user
    /// callback, where it can be added to the user's CpuWallTiming using
    /// CpuWallTiming::add().
    template<typename F>
    class DeltaCpuWallTimer {
    public:
        explicit DeltaCpuWallTimer(F &&func) : func_(std::move(func)) {
        }

        ~DeltaCpuWallTimer() {
            func_(timer_.elapsed());
        }

    private:
        DeltaCpuWallTimeStopWatch timer_;
        F func_;
    };
} // namespace kumo::pollux
